import sys
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QTextEdit, QPushButton, QComboBox, 
                             QLabel, QMessageBox, QProgressBar, QGroupBox)
from PyQt6.QtCore import QThread, pyqtSignal, Qt
from PyQt6.QtGui import QFont, QInputMethodEvent

from ollama_client import OllamaClient
from config_manager import get_config_manager
from config_window import ConfigWindow


class TranslationThread(QThread):
    finished = pyqtSignal(str)
    error = pyqtSignal(str)
    
    def __init__(self, client, text, source_lang, target_lang, model, custom_prompt=""):
        super().__init__()
        self.client = client
        self.text = text
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.model = model
        self.custom_prompt = custom_prompt
    
    def run(self):
        try:
            result = self.client.translate(
                self.text, 
                self.model, 
                self.source_lang, 
                self.target_lang,
                self.custom_prompt
            )
            if result:
                self.finished.emit(result)
            else:
                self.error.emit("Translation failed or returned empty result")
        except Exception as e:
            self.error.emit(str(e))


class StreamingTranslationThread(QThread):
    update = pyqtSignal(str)  # Emit partial translation
    finished = pyqtSignal()   # Signal when translation is complete
    error = pyqtSignal(str)
    
    def __init__(self, client, text, source_lang, target_lang, model, custom_prompt=""):
        super().__init__()
        self.client = client
        self.text = text
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.model = model
        self.custom_prompt = custom_prompt
        self._stop_requested = False
    
    def stop(self):
        """Request to stop the translation."""
        self._stop_requested = True
    
    def run(self):
        try:
            translation_buffer = ""
            
            for chunk in self.client.translate_stream(
                self.text, 
                self.model, 
                self.source_lang, 
                self.target_lang,
                self.custom_prompt
            ):
                if self._stop_requested:
                    break
                    
                translation_buffer += chunk
                # Emit the current translation state
                self.update.emit(translation_buffer)
                
            if not self._stop_requested:
                self.finished.emit()
                
        except Exception as e:
            if not self._stop_requested:
                self.error.emit(str(e))


class TranslatorGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.client = OllamaClient()
        self.config_manager = get_config_manager()
        self.available_models = []
        self.current_preset = None
        self.init_ui()
        self.load_presets()
        self.check_connection()
        self.apply_theme()
    
    def init_ui(self):
        self.setWindowTitle("Ollama Translator")
        self.setGeometry(100, 100, 900, 500)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        layout = QVBoxLayout(central_widget)
        
        # Simple header with connection status and config button
        header_layout = QHBoxLayout()
        
        # Connection status (minimal)
        self.connection_status_label = QLabel("🔍 Checking connection...")
        header_layout.addWidget(self.connection_status_label)
        
        header_layout.addStretch()
        
        # Config button
        self.config_btn = QPushButton("⚙️")
        self.config_btn.setToolTip("Configuration")
        self.config_btn.setMaximumWidth(40)
        self.config_btn.clicked.connect(self.open_config)
        header_layout.addWidget(self.config_btn)
        
        layout.addLayout(header_layout)
        
        # Simple language selection
        lang_layout = QHBoxLayout()
        
        lang_layout.addWidget(QLabel("From:"))
        self.source_lang_combo = QComboBox()
        self.source_lang_combo.addItems(["English", "Chinese", "French", "Spanish", "German", "Japanese"])
        lang_layout.addWidget(self.source_lang_combo)
        
        # Reverse button
        self.reverse_btn = QPushButton("⇄")
        self.reverse_btn.setToolTip("Reverse languages")
        self.reverse_btn.setMaximumWidth(40)
        self.reverse_btn.clicked.connect(self.reverse_languages)
        lang_layout.addWidget(self.reverse_btn)
        
        lang_layout.addWidget(QLabel("To:"))
        self.target_lang_combo = QComboBox()
        self.target_lang_combo.addItems(["Chinese", "English", "French", "Spanish", "German", "Japanese"])
        lang_layout.addWidget(self.target_lang_combo)
        
        lang_layout.addStretch()
        
        # Preset selection (minimal)
        preset_layout = QHBoxLayout()
        preset_layout.addWidget(QLabel("Style:"))
        self.preset_combo = QComboBox()
        self.preset_combo.setToolTip("Translation style")
        self.preset_combo.currentTextChanged.connect(self.on_preset_changed)
        preset_layout.addWidget(self.preset_combo)
        preset_layout.addStretch()
        
        layout.addLayout(lang_layout)
        layout.addLayout(preset_layout)
        
        # Text areas (larger and cleaner)
        text_layout = QHBoxLayout()
        
        # Source text
        self.source_text = QTextEdit()
        self.source_text.setPlaceholderText("Enter text to translate...")
        self.source_text.setAttribute(Qt.WidgetAttribute.WA_InputMethodEnabled, True)
        text_layout.addWidget(self.source_text)
        
        # Target text
        self.target_text = QTextEdit()
        self.target_text.setReadOnly(True)
        self.target_text.setPlaceholderText("Translation will appear here...")
        self.target_text.setAttribute(Qt.WidgetAttribute.WA_InputMethodEnabled, True)
        text_layout.addWidget(self.target_text)
        
        layout.addLayout(text_layout)
        
        # Progress and status (minimal)
        status_layout = QHBoxLayout()
        
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.progress_bar.setMaximumHeight(6)
        status_layout.addWidget(self.progress_bar)
        
        self.status_label = QLabel("Ready")
        status_layout.addWidget(self.status_label)
        
        layout.addLayout(status_layout)
        
        # Action buttons (clean and simple)
        button_layout = QHBoxLayout()
        
        self.translate_btn = QPushButton("Translate")
        self.translate_btn.clicked.connect(self.translate_text)
        button_layout.addWidget(self.translate_btn)
        
        self.cancel_btn = QPushButton("Cancel")
        self.cancel_btn.clicked.connect(self.cancel_translation)
        self.cancel_btn.setVisible(False)
        button_layout.addWidget(self.cancel_btn)
        
        self.clear_btn = QPushButton("Clear")
        self.clear_btn.clicked.connect(self.clear_text)
        button_layout.addWidget(self.clear_btn)
        
        button_layout.addStretch()
        layout.addLayout(button_layout)
        
        # Apply font settings
        self.apply_font_settings()
    
    def check_connection(self):
        """Check Ollama connection and update UI accordingly."""
        success, message = self.client.test_connection()
        
        if success:
            self.connection_status_label.setText("✓ Connected to Ollama")
            self.connection_status_label.setStyleSheet("color: green;")
            self.load_models()
        else:
            self.connection_status_label.setText("✗ " + message)
            self.connection_status_label.setStyleSheet("color: red;")
            self.translate_btn.setEnabled(False)
    
    def load_presets(self):
        """Load available presets."""
        self.preset_combo.clear()
        self.preset_combo.addItem("None (Custom)", None)
        
        presets = self.config_manager.presets
        for preset in presets:
            if preset.enabled:
                display_name = f"{preset.name} - {preset.description}"
                self.preset_combo.addItem(display_name, preset)
    
    def load_models(self):
        """Load available models from Ollama."""
        self.available_models = self.client.get_models()
        
        # Model selection is now in configuration window
        if not self.available_models:
            self.translate_btn.setEnabled(False)
    
    def on_preset_changed(self, preset_name):
        """Handle preset selection change."""
        preset_data = self.preset_combo.currentData()
        if preset_data:
            self.current_preset = preset_data
            # Apply preset settings
            self.source_lang_combo.setCurrentText(self.current_preset.source_language)
            self.target_lang_combo.setCurrentText(self.current_preset.target_language)
        else:
            self.current_preset = None
    
    def open_settings(self):
        """Open settings dialog."""
        dialog = SettingsDialog(self)
        dialog.exec()
        # Reload presets after settings dialog closes
        self.load_presets()
    
    def translate_text(self):
        """Start translation process with streaming support."""
        text = self.source_text.toPlainText().strip()
        if not text:
            QMessageBox.warning(self, "Warning", "Please enter text to translate")
            return
        
        if not self.available_models:
            QMessageBox.warning(self, "Warning", "No models available for translation")
            return
        
        source_lang = self.source_lang_combo.currentText()
        target_lang = self.target_lang_combo.currentText()
        # Get model from configuration
        model = self.config_manager.get_selected_model()
        
        if source_lang == target_lang:
            QMessageBox.warning(self, "Warning", "Source and target languages cannot be the same")
            return
        
        # Get custom prompt from preset if selected
        custom_prompt = ""
        if self.current_preset:
            custom_prompt = self.current_preset.custom_prompt
        
        # Clear target text and prepare for new translation
        self.target_text.clear()
        
        # Disable UI during translation
        self.translate_btn.setEnabled(False)
        self.cancel_btn.setVisible(True)
        self.progress_bar.setVisible(True)
        self.progress_bar.setRange(0, 0)  # Indeterminate progress
        self.status_label.setText("Translating...")
        
        # Choose between streaming and regular translation
        if self.config_manager.get_streaming_enabled():
            # Use streaming translation
            self.translation_thread = StreamingTranslationThread(
                self.client, text, source_lang, target_lang, model, custom_prompt
            )
            self.translation_thread.update.connect(self.streaming_update)
            self.translation_thread.finished.connect(self.translation_finished)
            self.translation_thread.error.connect(self.translation_error)
        else:
            # Use regular translation
            self.translation_thread = TranslationThread(
                self.client, text, source_lang, target_lang, model, custom_prompt
            )
            self.translation_thread.finished.connect(self.translation_finished)
            self.translation_thread.error.connect(self.translation_error)
        
        self.translation_thread.start()
    
    def streaming_update(self, partial_translation):
        """Update UI with partial translation during streaming."""
        self.target_text.setPlainText(partial_translation)
        # Auto-scroll to bottom to show latest content
        cursor = self.target_text.textCursor()
        cursor.movePosition(cursor.MoveOperation.End)
        self.target_text.setTextCursor(cursor)
    
    def translation_finished(self, result=None):
        """Handle successful translation completion."""
        if result is not None:
            # For non-streaming translation
            self.target_text.setPlainText(result)
        
        self.progress_bar.setVisible(False)
        self.translate_btn.setEnabled(True)
        self.cancel_btn.setVisible(False)
        self.status_label.setText("Translation completed")
    
    def translation_error(self, error_message):
        """Handle translation errors."""
        QMessageBox.critical(self, "Translation Error", error_message)
        self.progress_bar.setVisible(False)
        self.translate_btn.setEnabled(True)
        self.cancel_btn.setVisible(False)
        self.status_label.setText("Translation failed")
    
    def cancel_translation(self):
        """Cancel ongoing translation."""
        if hasattr(self, 'translation_thread') and self.translation_thread.isRunning():
            # Stop the streaming thread
            if isinstance(self.translation_thread, StreamingTranslationThread):
                self.translation_thread.stop()
            
            self.translation_thread.terminate()
            self.translation_thread.wait()
            
            self.progress_bar.setVisible(False)
            self.translate_btn.setEnabled(True)
            self.cancel_btn.setVisible(False)
            self.status_label.setText("Translation cancelled")
    
    def apply_theme(self):
        """Apply current theme to the GUI."""
        theme = self.config_manager.get_current_theme()
        
        # Main window stylesheet
        stylesheet = f"""
            QMainWindow {{
                background-color: {theme['background']};
                color: {theme['foreground']};
            }}
            
            QWidget {{
                background-color: {theme['background']};
                color: {theme['foreground']};
            }}
            
            QGroupBox {{
                background-color: {theme['group_bg']};
                color: {theme['foreground']};
                border: 2px solid {theme['group_border']};
                border-radius: 5px;
                margin-top: 1ex;
                padding-top: 10px;
                font-weight: bold;
            }}
            
            QGroupBox::title {{
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px 0 5px;
            }}
            
            QTextEdit {{
                background-color: {theme['text_edit_bg']};
                color: {theme['text_edit_fg']};
                border: 1px solid {theme['group_border']};
                border-radius: 3px;
                padding: 5px;
            }}
            
            QPushButton {{
                background-color: {theme['button_bg']};
                color: {theme['button_fg']};
                border: none;
                border-radius: 3px;
                padding: 5px 10px;
                font-weight: bold;
            }}
            
            QPushButton:hover {{
                background-color: {theme['primary']};
            }}
            
            QPushButton:pressed {{
                background-color: {theme['accent']};
            }}
            
            QPushButton:disabled {{
                background-color: {theme['secondary']};
                color: gray;
            }}
            
            QComboBox {{
                background-color: {theme['text_edit_bg']};
                color: {theme['text_edit_fg']};
                border: 1px solid {theme['group_border']};
                border-radius: 3px;
                padding: 3px;
            }}
            
            QComboBox::drop-down {{
                border: none;
            }}
            
            QComboBox QAbstractItemView {{
                background-color: {theme['text_edit_bg']};
                color: {theme['text_edit_fg']};
                selection-background-color: {theme['primary']};
            }}
            
            QCheckBox {{
                color: {theme['foreground']};
            }}
            
            QCheckBox::indicator {{
                width: 13px;
                height: 13px;
            }}
            
            QCheckBox::indicator:unchecked {{
                border: 1px solid {theme['group_border']};
                background-color: {theme['text_edit_bg']};
            }}
            
            QCheckBox::indicator:checked {{
                border: 1px solid {theme['primary']};
                background-color: {theme['primary']};
            }}
            
            QProgressBar {{
                border: 1px solid {theme['group_border']};
                border-radius: 3px;
                text-align: center;
                background-color: {theme['secondary']};
            }}
            
            QProgressBar::chunk {{
                background-color: {theme['primary']};
                border-radius: 2px;
            }}
            
            QLabel {{
                color: {theme['foreground']};
            }}
        """
        
        self.setStyleSheet(stylesheet)
    
    def open_settings(self):
        """Open settings dialog."""
        dialog = SettingsDialog(self)
        result = dialog.exec()
        # Reload presets after settings dialog closes
        self.load_presets()
        # Reapply theme if it was changed
        self.apply_theme()
    
    def reverse_languages(self):
        """Reverse source and target languages and swap text content."""
        # Get current selections
        source_lang = self.source_lang_combo.currentText()
        target_lang = self.target_lang_combo.currentText()
        
        # Don't reverse if languages are the same
        if source_lang == target_lang:
            return
        
        # Swap the language selections
        self.source_lang_combo.setCurrentText(target_lang)
        self.target_lang_combo.setCurrentText(source_lang)
        
        # Get current text content
        source_text = self.source_text.toPlainText()
        target_text = self.target_text.toPlainText()
        
        # Swap the text content only if target text is not empty
        if target_text.strip():
            self.source_text.setPlainText(target_text)
            self.target_text.setPlainText(source_text)
    
    def clear_text(self):
        """Clear both text areas."""
        self.source_text.clear()
        self.target_text.clear()
    
    def apply_font_settings(self):
        """Apply font settings from configuration."""
        font_family = self.config_manager.get_font_family()
        font_size = self.config_manager.get_font_size()
        
        if font_family == "System Default":
            font = QFont()
        else:
            font = QFont(font_family)
        
        font.setPointSize(font_size)
        
        # Apply to text areas
        self.source_text.setFont(font)
        self.target_text.setFont(font)
    
    def open_config(self):
        """Open configuration window."""
        dialog = ConfigWindow(self)
        dialog.config_changed.connect(self.on_config_changed)
        dialog.exec()
    
    def on_config_changed(self):
        """Handle configuration changes."""
        # Reapply theme and font settings
        self.apply_theme()
        self.apply_font_settings()
        self.load_presets()
        
        # Update streaming setting
        self.load_models()


def main():
    app = QApplication(sys.argv)
    
    # Set application properties
    app.setApplicationName("Ollama Translator")
    app.setApplicationVersion("0.1.0")
    
    window = TranslatorGUI()
    window.show()
    
    sys.exit(app.exec())


if __name__ == "__main__":
    main()